import numpy as np
import pytorch_lightning as pl
import torch
from torch.utils.data import DataLoader

from train_img.dataloader_nuimages import NuImagesDataset, collate_pair_fn

# try:
#     from pretrain.dataloader_nuscenes import (
#         NuScenesMatchDataset,
#         minkunet_collate_pair_fn,
#     )
# except ImportError:
#     NuScenesMatchDataset = None
#     minkunet_collate_pair_fn = None
# try:
#     from pretrain.dataloader_nuscenes_spconv import NuScenesMatchDatasetSpconv, spconv_collate_pair_fn
# except ImportError:
#     NuScenesMatchDatasetSpconv = None
#     spconv_collate_pair_fn = None
# from utils.transforms import (
#     make_transforms_clouds,
#     make_transforms_asymmetrical,
#     make_transforms_asymmetrical_val,
# )


class TrainImgDataModule(pl.LightningDataModule):
    def __init__(self, config):
        super().__init__()
        self.config = config
        if config["num_gpus"]:
            self.batch_size = config["batch_size"] // config["num_gpus"]
        else:
            self.batch_size = config["batch_size"]

    def setup(self, stage):
        if self.config["dataset"].lower() == "nuimages":
            Dataset = NuImagesDataset
        else:
            raise Exception("Dataset Unknown")

        self.train_dataset = Dataset(
            phase="train",
            config=self.config,
        )
        print("Dataset Loaded")
        self.val_dataset = Dataset(phase="val", config=self.config)

    def train_dataloader(self):
        if self.config["num_gpus"]:
            num_workers = self.config["num_threads"] // self.config["num_gpus"]
        else:
            num_workers = self.config["num_threads"]
        default_collate_pair_fn = collate_pair_fn
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=num_workers,
            collate_fn=default_collate_pair_fn,
            pin_memory=True,
            drop_last=True,
            worker_init_fn=lambda id: np.random.seed(
                torch.initial_seed() // 2**32 + id
            ),
        )

    def val_dataloader(self):
        if self.config["num_gpus"]:
            num_workers = self.config["num_threads"] // self.config["num_gpus"]
        else:
            num_workers = self.config["num_threads"]
        default_collate_pair_fn = collate_pair_fn
        return DataLoader(
            self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=num_workers,
            collate_fn=default_collate_pair_fn,
            pin_memory=True,
            drop_last=False,
            worker_init_fn=lambda id: np.random.seed(
                torch.initial_seed() // 2**32 + id
            ),
        )
